"""
异步任务处理服务

提供消息队列系统和长时间任务异步处理，包括：
- 任务队列管理
- 异步任务执行
- 任务调度优化
- 任务状态监控
"""
import asyncio
import json
import uuid
from typing import Any, Dict, List, Optional, Callable, Union
from datetime import datetime, timedelta
from dataclasses import dataclass, field
from enum import Enum
from collections import defaultdict, deque
import time
import traceback

try:
    import aioredis
    REDIS_AVAILABLE = True
except ImportError:
    REDIS_AVAILABLE = False


class TaskStatus(Enum):
    """任务状态枚举"""
    PENDING = "pending"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELLED = "cancelled"
    RETRYING = "retrying"


class TaskPriority(Enum):
    """任务优先级枚举"""
    LOW = 1
    NORMAL = 2
    HIGH = 3
    URGENT = 4
    CRITICAL = 5


@dataclass
class TaskResult:
    """任务结果"""
    task_id: str
    status: TaskStatus
    result: Any = None
    error: Optional[str] = None
    execution_time: float = 0.0
    completed_at: Optional[datetime] = None
    metadata: Dict[str, Any] = field(default_factory=dict)
@
dataclass
class Task:
    """任务定义"""
    task_id: str = field(default_factory=lambda: str(uuid.uuid4()))
    name: str = ""
    function: Optional[Callable] = None
    args: tuple = field(default_factory=tuple)
    kwargs: Dict[str, Any] = field(default_factory=dict)
    priority: TaskPriority = TaskPriority.NORMAL
    status: TaskStatus = TaskStatus.PENDING
    created_at: datetime = field(default_factory=datetime.now)
    scheduled_at: Optional[datetime] = None
    started_at: Optional[datetime] = None
    completed_at: Optional[datetime] = None
    max_retries: int = 3
    retry_count: int = 0
    timeout_seconds: Optional[int] = None
    dependencies: List[str] = field(default_factory=list)
    tags: List[str] = field(default_factory=list)
    metadata: Dict[str, Any] = field(default_factory=dict)
    result: Optional[TaskResult] = None


class TaskQueue:
    """任务队列"""
    
    def __init__(self, name: str, max_size: int = 1000):
        self.name = name
        self.max_size = max_size
        self.tasks: deque = deque()
        self.priority_queues: Dict[TaskPriority, deque] = {
            priority: deque() for priority in TaskPriority
        }
        self.task_lookup: Dict[str, Task] = {}
        self.stats = {
            'total_added': 0,
            'total_processed': 0,
            'current_size': 0,
            'avg_wait_time': 0.0
        }
    
    async def add_task(self, task: Task) -> bool:
        """添加任务到队列"""
        if len(self.tasks) >= self.max_size:
            return False
        
        # 添加到优先级队列
        self.priority_queues[task.priority].append(task)
        self.tasks.append(task)
        self.task_lookup[task.task_id] = task
        
        self.stats['total_added'] += 1
        self.stats['current_size'] = len(self.tasks)
        
        return True
    
    async def get_next_task(self) -> Optional[Task]:
        """获取下一个任务（按优先级）"""
        # 按优先级从高到低检查
        for priority in sorted(TaskPriority, key=lambda x: x.value, reverse=True):
            if self.priority_queues[priority]:
                task = self.priority_queues[priority].popleft()
                self.tasks.remove(task)
                self.stats['current_size'] = len(self.tasks)
                return task
        
        return None
    
    async def remove_task(self, task_id: str) -> bool:
        """移除任务"""
        if task_id in self.task_lookup:
            task = self.task_lookup[task_id]
            
            # 从各个队列中移除
            if task in self.tasks:
                self.tasks.remove(task)
            
            if task in self.priority_queues[task.priority]:
                self.priority_queues[task.priority].remove(task)
            
            del self.task_lookup[task_id]
            self.stats['current_size'] = len(self.tasks)
            
            return True
        
        return False
    
    def get_task(self, task_id: str) -> Optional[Task]:
        """获取任务"""
        return self.task_lookup.get(task_id)
    
    def get_stats(self) -> Dict[str, Any]:
        """获取队列统计"""
        return {
            'name': self.name,
            'current_size': len(self.tasks),
            'max_size': self.max_size,
            'priority_distribution': {
                priority.name: len(queue) 
                for priority, queue in self.priority_queues.items()
            },
            'stats': self.stats
        }


class TaskWorker:
    """任务工作者"""
    
    def __init__(self, worker_id: str, task_queue: TaskQueue):
        self.worker_id = worker_id
        self.task_queue = task_queue
        self.current_task: Optional[Task] = None
        self.is_running = False
        self.stats = {
            'tasks_processed': 0,
            'tasks_failed': 0,
            'total_execution_time': 0.0,
            'avg_execution_time': 0.0,
            'last_activity': None
        }
    
    async def start(self):
        """启动工作者"""
        self.is_running = True
        print(f"任务工作者 {self.worker_id} 已启动")
        
        while self.is_running:
            try:
                # 获取下一个任务
                task = await self.task_queue.get_next_task()
                
                if task is None:
                    # 没有任务时短暂休眠
                    await asyncio.sleep(0.1)
                    continue
                
                # 执行任务
                await self._execute_task(task)
                
            except Exception as e:
                print(f"工作者 {self.worker_id} 执行任务时出错: {e}")
                await asyncio.sleep(1)
    
    async def stop(self):
        """停止工作者"""
        self.is_running = False
        print(f"任务工作者 {self.worker_id} 已停止")
    
    async def _execute_task(self, task: Task):
        """执行任务"""
        self.current_task = task
        task.status = TaskStatus.RUNNING
        task.started_at = datetime.now()
        
        start_time = time.time()
        
        try:
            # 检查依赖
            if not await self._check_dependencies(task):
                task.status = TaskStatus.PENDING
                # 重新加入队列
                await self.task_queue.add_task(task)
                return
            
            # 执行任务函数
            if asyncio.iscoroutinefunction(task.function):
                result = await asyncio.wait_for(
                    task.function(*task.args, **task.kwargs),
                    timeout=task.timeout_seconds
                )
            else:
                result = task.function(*task.args, **task.kwargs)
            
            # 任务成功完成
            execution_time = time.time() - start_time
            task.status = TaskStatus.COMPLETED
            task.completed_at = datetime.now()
            task.result = TaskResult(
                task_id=task.task_id,
                status=TaskStatus.COMPLETED,
                result=result,
                execution_time=execution_time,
                completed_at=task.completed_at
            )
            
            # 更新统计
            self.stats['tasks_processed'] += 1
            self.stats['total_execution_time'] += execution_time
            self.stats['avg_execution_time'] = (
                self.stats['total_execution_time'] / self.stats['tasks_processed']
            )
            self.stats['last_activity'] = datetime.now()
            
            print(f"任务 {task.task_id} 执行成功 (耗时: {execution_time:.2f}s)")
            
        except asyncio.TimeoutError:
            # 任务超时
            task.status = TaskStatus.FAILED
            task.completed_at = datetime.now()
            task.result = TaskResult(
                task_id=task.task_id,
                status=TaskStatus.FAILED,
                error="任务执行超时",
                execution_time=time.time() - start_time,
                completed_at=task.completed_at
            )
            
            await self._handle_task_failure(task)
            
        except Exception as e:
            # 任务执行失败
            execution_time = time.time() - start_time
            task.status = TaskStatus.FAILED
            task.completed_at = datetime.now()
            task.result = TaskResult(
                task_id=task.task_id,
                status=TaskStatus.FAILED,
                error=str(e),
                execution_time=execution_time,
                completed_at=task.completed_at,
                metadata={'traceback': traceback.format_exc()}
            )
            
            await self._handle_task_failure(task)
            
        finally:
            self.current_task = None
    
    async def _check_dependencies(self, task: Task) -> bool:
        """检查任务依赖"""
        if not task.dependencies:
            return True
        
        # 这里应该检查依赖任务的状态
        # 简化实现：假设依赖检查总是通过
        return True
    
    async def _handle_task_failure(self, task: Task):
        """处理任务失败"""
        self.stats['tasks_failed'] += 1
        
        # 检查是否需要重试
        if task.retry_count < task.max_retries:
            task.retry_count += 1
            task.status = TaskStatus.RETRYING
            
            # 重新加入队列
            await self.task_queue.add_task(task)
            print(f"任务 {task.task_id} 重试 ({task.retry_count}/{task.max_retries})")
        else:
            print(f"任务 {task.task_id} 最终失败: {task.result.error}")
    
    def get_stats(self) -> Dict[str, Any]:
        """获取工作者统计"""
        return {
            'worker_id': self.worker_id,
            'is_running': self.is_running,
            'current_task': self.current_task.task_id if self.current_task else None,
            'stats': self.stats
        }class A
syncTaskService:
    """异步任务服务"""
    
    def __init__(self, max_workers: int = 4, redis_url: Optional[str] = None):
        self.max_workers = max_workers
        self.redis_url = redis_url
        
        # 任务队列
        self.task_queues: Dict[str, TaskQueue] = {}
        self.default_queue = "default"
        
        # 工作者
        self.workers: List[TaskWorker] = []
        self.worker_tasks: List[asyncio.Task] = []
        
        # 任务存储
        self.task_storage: Dict[str, Task] = {}
        self.task_results: Dict[str, TaskResult] = {}
        
        # 调度器
        self.scheduler_running = False
        self.scheduler_task: Optional[asyncio.Task] = None
        
        # 统计信息
        self.service_stats = {
            'total_tasks_submitted': 0,
            'total_tasks_completed': 0,
            'total_tasks_failed': 0,
            'service_start_time': None,
            'uptime_seconds': 0
        }
        
        # Redis客户端（如果可用）
        self.redis_client: Optional[aioredis.Redis] = None
        
        # 初始化默认队列
        self._create_queue(self.default_queue)
    
    async def initialize(self):
        """初始化服务"""
        # 尝试连接Redis
        if self.redis_url and REDIS_AVAILABLE:
            try:
                self.redis_client = aioredis.from_url(self.redis_url)
                await self.redis_client.ping()
                print("✅ Redis任务存储已连接")
            except Exception as e:
                print(f"⚠️ Redis连接失败，使用内存存储: {e}")
        
        # 启动工作者
        await self._start_workers()
        
        # 启动调度器
        await self._start_scheduler()
        
        self.service_stats['service_start_time'] = datetime.now()
        print(f"🚀 异步任务服务已启动 ({self.max_workers} 个工作者)")
    
    def _create_queue(self, queue_name: str, max_size: int = 1000):
        """创建任务队列"""
        if queue_name not in self.task_queues:
            self.task_queues[queue_name] = TaskQueue(queue_name, max_size)
    
    async def _start_workers(self):
        """启动工作者"""
        default_queue = self.task_queues[self.default_queue]
        
        for i in range(self.max_workers):
            worker_id = f"worker_{i}"
            worker = TaskWorker(worker_id, default_queue)
            self.workers.append(worker)
            
            # 启动工作者任务
            worker_task = asyncio.create_task(worker.start())
            self.worker_tasks.append(worker_task)
    
    async def _start_scheduler(self):
        """启动任务调度器"""
        self.scheduler_running = True
        self.scheduler_task = asyncio.create_task(self._scheduler_loop())
    
    async def _scheduler_loop(self):
        """调度器循环"""
        while self.scheduler_running:
            try:
                # 处理定时任务
                await self._process_scheduled_tasks()
                
                # 更新统计信息
                await self._update_service_stats()
                
                # 清理完成的任务
                await self._cleanup_completed_tasks()
                
                # 休眠
                await asyncio.sleep(1)
                
            except Exception as e:
                print(f"调度器错误: {e}")
                await asyncio.sleep(5)
    
    async def submit_task(self, function: Callable, *args, 
                         name: str = "", priority: TaskPriority = TaskPriority.NORMAL,
                         queue_name: str = None, timeout_seconds: Optional[int] = None,
                         max_retries: int = 3, scheduled_at: Optional[datetime] = None,
                         dependencies: List[str] = None, tags: List[str] = None,
                         **kwargs) -> str:
        """提交任务"""
        # 创建任务
        task = Task(
            name=name or function.__name__,
            function=function,
            args=args,
            kwargs=kwargs,
            priority=priority,
            timeout_seconds=timeout_seconds,
            max_retries=max_retries,
            scheduled_at=scheduled_at,
            dependencies=dependencies or [],
            tags=tags or []
        )
        
        # 选择队列
        target_queue = queue_name or self.default_queue
        if target_queue not in self.task_queues:
            self._create_queue(target_queue)
        
        # 存储任务
        self.task_storage[task.task_id] = task
        
        # 如果是定时任务，暂不加入队列
        if scheduled_at and scheduled_at > datetime.now():
            print(f"任务 {task.task_id} 已安排在 {scheduled_at} 执行")
        else:
            # 立即加入队列
            success = await self.task_queues[target_queue].add_task(task)
            if not success:
                raise Exception(f"队列 {target_queue} 已满，无法添加任务")
        
        self.service_stats['total_tasks_submitted'] += 1
        
        # 持久化到Redis（如果可用）
        if self.redis_client:
            await self._persist_task(task)
        
        print(f"任务 {task.task_id} ({name}) 已提交")
        return task.task_id
    
    async def get_task_status(self, task_id: str) -> Optional[Dict[str, Any]]:
        """获取任务状态"""
        if task_id in self.task_storage:
            task = self.task_storage[task_id]
            return {
                'task_id': task.task_id,
                'name': task.name,
                'status': task.status.value,
                'created_at': task.created_at.isoformat(),
                'started_at': task.started_at.isoformat() if task.started_at else None,
                'completed_at': task.completed_at.isoformat() if task.completed_at else None,
                'retry_count': task.retry_count,
                'max_retries': task.max_retries,
                'result': task.result.__dict__ if task.result else None
            }
        
        return None
    
    async def get_task_result(self, task_id: str, timeout: Optional[float] = None) -> Optional[TaskResult]:
        """获取任务结果（阻塞等待）"""
        start_time = time.time()
        
        while True:
            if task_id in self.task_storage:
                task = self.task_storage[task_id]
                
                if task.result:
                    return task.result
                
                if task.status in [TaskStatus.FAILED, TaskStatus.CANCELLED]:
                    return task.result
            
            # 检查超时
            if timeout and (time.time() - start_time) > timeout:
                return None
            
            await asyncio.sleep(0.1)
    
    async def cancel_task(self, task_id: str) -> bool:
        """取消任务"""
        if task_id in self.task_storage:
            task = self.task_storage[task_id]
            
            if task.status == TaskStatus.PENDING:
                # 从队列中移除
                for queue in self.task_queues.values():
                    if await queue.remove_task(task_id):
                        task.status = TaskStatus.CANCELLED
                        task.completed_at = datetime.now()
                        print(f"任务 {task_id} 已取消")
                        return True
            elif task.status == TaskStatus.RUNNING:
                # 正在运行的任务标记为取消（实际取消需要任务函数支持）
                task.status = TaskStatus.CANCELLED
                print(f"任务 {task_id} 已标记为取消")
                return True
        
        return False
    
    async def _process_scheduled_tasks(self):
        """处理定时任务"""
        current_time = datetime.now()
        
        for task_id, task in list(self.task_storage.items()):
            if (task.scheduled_at and 
                task.scheduled_at <= current_time and 
                task.status == TaskStatus.PENDING):
                
                # 将定时任务加入队列
                queue = self.task_queues[self.default_queue]
                await queue.add_task(task)
                task.scheduled_at = None  # 清除定时标记
    
    async def _update_service_stats(self):
        """更新服务统计"""
        if self.service_stats['service_start_time']:
            uptime = datetime.now() - self.service_stats['service_start_time']
            self.service_stats['uptime_seconds'] = uptime.total_seconds()
        
        # 统计完成和失败的任务
        completed_count = 0
        failed_count = 0
        
        for task in self.task_storage.values():
            if task.status == TaskStatus.COMPLETED:
                completed_count += 1
            elif task.status == TaskStatus.FAILED:
                failed_count += 1
        
        self.service_stats['total_tasks_completed'] = completed_count
        self.service_stats['total_tasks_failed'] = failed_count
    
    async def _cleanup_completed_tasks(self):
        """清理完成的任务"""
        # 清理超过24小时的已完成任务
        cutoff_time = datetime.now() - timedelta(hours=24)
        tasks_to_remove = []
        
        for task_id, task in self.task_storage.items():
            if (task.completed_at and 
                task.completed_at < cutoff_time and 
                task.status in [TaskStatus.COMPLETED, TaskStatus.FAILED, TaskStatus.CANCELLED]):
                tasks_to_remove.append(task_id)
        
        for task_id in tasks_to_remove:
            del self.task_storage[task_id]
            if task_id in self.task_results:
                del self.task_results[task_id]
    
    async def _persist_task(self, task: Task):
        """持久化任务到Redis"""
        if self.redis_client:
            try:
                task_data = {
                    'task_id': task.task_id,
                    'name': task.name,
                    'status': task.status.value,
                    'created_at': task.created_at.isoformat(),
                    'priority': task.priority.value,
                    'max_retries': task.max_retries,
                    'retry_count': task.retry_count,
                    'tags': task.tags,
                    'metadata': task.metadata
                }
                
                await self.redis_client.hset(
                    f"task:{task.task_id}",
                    mapping=task_data
                )
                
                # 设置过期时间（7天）
                await self.redis_client.expire(f"task:{task.task_id}", 604800)
                
            except Exception as e:
                print(f"持久化任务失败 {task.task_id}: {e}")
    
    async def get_service_stats(self) -> Dict[str, Any]:
        """获取服务统计"""
        await self._update_service_stats()
        
        # 获取队列统计
        queue_stats = {}
        for queue_name, queue in self.task_queues.items():
            queue_stats[queue_name] = queue.get_stats()
        
        # 获取工作者统计
        worker_stats = []
        for worker in self.workers:
            worker_stats.append(worker.get_stats())
        
        return {
            'service_stats': self.service_stats,
            'queue_stats': queue_stats,
            'worker_stats': worker_stats,
            'total_queues': len(self.task_queues),
            'total_workers': len(self.workers),
            'active_tasks': len([t for t in self.task_storage.values() 
                               if t.status == TaskStatus.RUNNING])
        }
    
    async def shutdown(self):
        """关闭服务"""
        print("🛑 正在关闭异步任务服务...")
        
        # 停止调度器
        self.scheduler_running = False
        if self.scheduler_task:
            self.scheduler_task.cancel()
            try:
                await self.scheduler_task
            except asyncio.CancelledError:
                pass
        
        # 停止工作者
        for worker in self.workers:
            await worker.stop()
        
        # 等待工作者任务完成
        for worker_task in self.worker_tasks:
            worker_task.cancel()
            try:
                await worker_task
            except asyncio.CancelledError:
                pass
        
        # 关闭Redis连接
        if self.redis_client:
            await self.redis_client.close()
        
        print("✅ 异步任务服务已关闭")


# 便捷装饰器
def async_task(service: AsyncTaskService, priority: TaskPriority = TaskPriority.NORMAL,
               queue_name: str = None, timeout_seconds: Optional[int] = None,
               max_retries: int = 3):
    """异步任务装饰器"""
    def decorator(func):
        async def wrapper(*args, **kwargs):
            task_id = await service.submit_task(
                func, *args,
                name=func.__name__,
                priority=priority,
                queue_name=queue_name,
                timeout_seconds=timeout_seconds,
                max_retries=max_retries,
                **kwargs
            )
            return task_id
        
        wrapper.get_result = lambda task_id, timeout=None: service.get_task_result(task_id, timeout)
        wrapper.get_status = lambda task_id: service.get_task_status(task_id)
        
        return wrapper
    return decorator